#include <numeric>

#include "pub_sub_common.h"

std::chrono::high_resolution_clock::time_point start_time;
std::chrono::high_resolution_clock::time_point end_time;

void subscriber(std::vector<std::chrono::duration<double, std::micro>>& latencies) {
  int sock = nn_socket(AF_SP, NN_SUB);
  if (sock < 0) {
    fprintf(stderr, "Error creating subscriber socket: %s\n", nn_strerror(nn_errno()));
    return;
  }
  if (nn_connect(sock, URL) < 0) {
    fprintf(stderr, "Error connecting subscriber socket: %s\n", nn_strerror(nn_errno()));
    nn_close(sock);
    return;
  }
  nn_setsockopt(sock, NN_SUB, NN_SUB_SUBSCRIBE, "", 0);
  int recv_buf_len{0};
  size_t sz = sizeof(recv_buf_len);
  nn_getsockopt(sock, NN_SOL_SOCKET, NN_RCVBUF, &recv_buf_len, &sz);
  fprintf(stdout, "### Default NN_RECVBUF size %d.\n", recv_buf_len);

#if ENABLE_CHANGE_BUFSIZE
  // Set receive buffer size
  nn_setsockopt(sock, NN_SOL_SOCKET, NN_RCVBUF, &BUFFER_SIZE, sizeof(BUFFER_SIZE));
  nn_getsockopt(sock, NN_SOL_SOCKET, NN_RCVBUF, &recv_buf_len, &sz);
  fprintf(stdout, "### After set NN_RECVBUF size %d.\n", recv_buf_len);
#endif

  // Set timeout
  nn_setsockopt(sock, NN_SOL_SOCKET, NN_RCVTIMEO, &RECEIVE_TIMEOUT, sizeof(RECEIVE_TIMEOUT));
  nn_setsockopt(sock, NN_SOL_SOCKET, NN_RCVMAXSIZE, &RECVMAXSIZE, sizeof(RECVMAXSIZE));

  fprintf(stdout, "Subscriber socket connected to %s\n", URL);
  FileLock lock;
  lock.lock();
  char* buf = nullptr;
  nn_pollfd pfd[1];
  pfd[0].fd = sock;
  pfd[0].events = NN_POLLIN;
  uint32_t recv_count{0};

  fprintf(stdout, "@@ Start to recv, current time: %s\n", get_currenttime_millis().c_str());
  start_time = std::chrono::high_resolution_clock::now();
  while (recv_count < NUM_MESSAGES) {
    int rc = nn_poll(pfd, 1, POLL_TIMEOUT);
    if (rc < 0) {
      fprintf(stderr, "Error polling socket: %s\n", nn_strerror(nn_errno()));
      return;
    } else if (rc > 0) {
      if ((pfd[0].revents & NN_POLLIN) != 0U) {
        int bytes = nn_recv(sock, &buf, NN_MSG, 0);
        if (bytes < 0) {
          fprintf(stderr, "Error receiving message %d: %s\n", recv_count + 1, nn_strerror(nn_errno()));
          if (buf) {
            nn_freemsg(buf);
          }
          return;
        }
#if DEBUG_ON
        fprintf(stdout, "Received message %d of size %d bytes\n", recv_count + 1, bytes);
#endif
#if ENABLE_CRC_CHECK
        // Process received message
        uint32_t length;
        uint16_t crc;
        std::memcpy(&length, buf, sizeof(length));
        std::memcpy(&crc, buf + sizeof(length), sizeof(crc));

        uint16_t calculated_crc = calculate_crc2(buf + sizeof(length) + sizeof(crc), length);
        if (calculated_crc != crc) {
          fprintf(stderr, "CRC mismatch in message %d: expected %hu, got %hu\n", recv_count + 1, crc, calculated_crc);
        }
#if DEBUG_ON
        fprintf(stdout, "Received data block of size %u bytes with valid CRC in message %d\n", length, recv_count + 1);
#endif
#endif
        nn_freemsg(buf);
        recv_count += 1;
      }
    } else {
#if DEBUG_ON
      fprintf(stdout, "nn_poll timeout %d(ms)\n", POLL_TIMEOUT);
#endif
    }
  }
  end_time = std::chrono::high_resolution_clock::now();
  fprintf(stdout, "@@ Finish recv, current time: %s\n", get_currenttime_millis().c_str());
  nn_close(sock);
}

int main() {
  std::vector<std::chrono::duration<double, std::micro>> latencies;
  subscriber(latencies);
  auto total_latency = end_time - start_time;
  auto avg_latency = total_latency / NUM_MESSAGES;
  fprintf(stdout, "Subscriber receive %d messages.\nAverage latency: %lf milliseconds\n", NUM_MESSAGES,
          std::chrono::duration<double, std::milli>(avg_latency).count());
  return 0;
}
